import contextvars
import logging
import time
from typing import Optional

# 创建上下文变量来存储请求ID
request_id_var: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar('request_id', default=None)

class RequestContext:
    """
    请求上下文管理器，用于在异步调用链中传递请求ID
    """
    
    @staticmethod
    def set_request_id(request_id: str) -> None:
        """
        设置当前请求的ID
        
        参数:
        - request_id: 请求ID
        """
        request_id_var.set(request_id)
    
    @staticmethod
    def get_request_id() -> Optional[str]:
        """
        获取当前请求的ID
        
        返回:
        - 当前请求ID，如果没有设置则返回None
        """
        return request_id_var.get()
    
    @staticmethod
    def generate_request_id() -> str:
        """
        生成一个新的请求ID
        
        返回:
        - 基于时间戳的请求ID
        """
        return str(int(time.time() * 1000000))
    
    @staticmethod
    def clear_request_id() -> None:
        """
        清除当前请求ID
        """
        request_id_var.set(None)

class ContextLogger:
    """
    带上下文的日志记录器，自动添加请求ID前缀
    """
    
    @staticmethod
    def _format_message(message: str) -> str:
        """
        格式化日志消息，添加请求ID前缀
        
        参数:
        - message: 原始日志消息
        
        返回:
        - 格式化后的消息
        """
        request_id = RequestContext.get_request_id()
        if request_id:
            return f"[请求ID:{request_id}] {message}"
        return message
    
    @staticmethod
    def info(message: str, *args, **kwargs) -> None:
        """
        记录INFO级别日志，自动添加请求ID
        """
        formatted_message = ContextLogger._format_message(message)
        logging.info(formatted_message, *args, **kwargs)
    
    @staticmethod
    def error(message: str, *args, **kwargs) -> None:
        """
        记录ERROR级别日志，自动添加请求ID
        """
        formatted_message = ContextLogger._format_message(message)
        logging.error(formatted_message, *args, **kwargs)
    
    @staticmethod
    def warning(message: str, *args, **kwargs) -> None:
        """
        记录WARNING级别日志，自动添加请求ID
        """
        formatted_message = ContextLogger._format_message(message)
        logging.warning(formatted_message, *args, **kwargs)
    
    @staticmethod
    def debug(message: str, *args, **kwargs) -> None:
        """
        记录DEBUG级别日志，自动添加请求ID
        """
        formatted_message = ContextLogger._format_message(message)
        logging.debug(formatted_message, *args, **kwargs)

# 提供简便的导入别名
context_logger = ContextLogger()